import torch
from torch import nn
import torch.nn.functional as F

# 构造一个没有任何参数的自定义层,下面 的CenteredLayer类要从其输入中减去均值。要构建它，我们只需继承基础层类并实现前向传播功能
class CenteredLayer(nn.Module):
    def __init__(self):
        super(CenteredLayer, self).__init__()

    def forward(self, X):
        return X - X.mean()
# 提供数据测试一下
layer = CenteredLayer()
layer(torch.FloatTensor([1, 2, 3, 4, 5]))

# 现在，我们可以将层作为组件合并到更复杂的模型中。
net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())
Y = net(torch.rand(4, 8))
Y.mean()

# 下面我们继续定义具有参数的层，这些参数可以通过训练进行调整。我 们可以使用内置函数来创建参数，这些函数提供一些基本的管理功能。
# 实现自定义版本的全连接层。回想一下，该层需要两个参数，一个用于表示权重，另一个用于 表示偏置项。
# 在此实现中，我们使用修正线性单元作为激活函数。该层需要输入参数：in_units和units，分 别表示输入数和输出数。
class MyLinear(nn.Module):
    def __init__(self, in_units, units):
        super(MyLinear, self).__init__()
        self.weight = nn.Parameter(torch.randn(in_units, units))
        self.bias = nn.Parameter(torch.randn(units))
    def forward(self, X):
        H = torch.matmul(X, self.weight.data) + self.bias.data  # 注意这里的自定义实现要用weight和bias的data
        return F.relu(H)
# 实例化类进行测试
linear = MyLinear(5, 3)
linear.weight
linear(torch.rand(3, 5))

# 我们还可以使用自定义层构建模型，就像使用内置的全连接层一样使用自定义层
net = nn.Sequential(MyLinear(64, 8), MyLinear(8, 1))
net(torch.rand(3, 64))